/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
 * Description: 节点间AHC非对称拼接实现的all_reduce操作类实现文件
 */

 #include "all_reduce_ahc.h"

namespace hccl {
AllReduceAHC::AllReduceAHC(const HcclDispatcher dispatcher, const u64 reduceAttrBitMap,
    const u64 totalCount, const std::vector<std::vector<u32>> &subGroups) 
    : AHCExecutorBase(dispatcher, reduceAttrBitMap, totalCount, subGroups)
{
}

AllReduceAHC::~AllReduceAHC()
{
}

HcclResult AllReduceAHC::CommAHCInfoInit(std::vector<std::vector<u32>> &subGroups)
{
    commAHCBaseInfo_.reset(new (std::nothrow) CommAHCAlignInfo(subGroups));
    CHK_SMART_PTR_NULL(commAHCBaseInfo_);
    CHK_RET(commAHCBaseInfo_->Init());
    return HCCL_SUCCESS;
}

HcclResult AllReduceAHC::RunInterAllReduce(const u32 rank, const std::vector<LINK> &links, 
    const std::unique_ptr<CommAHCBaseInfo> &commAHCBaseInfo)
{
    HCCL_INFO("[AllReduceAHC][RunInterAllReduce] begin inter allreduce rank[%u]", rank);

    // 获取当前rank的组间rank
    u32 interRank = commAHCBaseInfo->GetInterRank(rank);

    // 创建执行算子实列
    std::unique_ptr<AlgTemplateBase> tempAlg;
    CHK_RET(commAHCBaseInfo->GetExecutorOpInstance(HcclCMDType::HCCL_CMD_ALLREDUCE, tempAlg, dispatcher_, reduceAttr_));

    std::vector<std::vector<Slice>> interSlicesVector;
    std::vector<std::vector<LINK>>  interLinksVector;
    std::vector<u32>                logicCardList;
    CHK_RET(commAHCBaseInfo->CalcInterSlicesAndLinks(rank, DataUnitSize(dataType_), count_, links,
        interLinksVector, interSlicesVector, logicCardList));

    if(interLinksVector.size() != interSlicesVector.size()) {
        HCCL_ERROR("[AllReduceAHC][RunInterAllReduce]rank[%u] linksVector size[%llu] is no equal to slicesVector size [%u]",
            rank, interLinksVector.size(), interSlicesVector.size());
        return HCCL_E_INTERNAL;
    }

    for (u32 i = 0; i < logicCardList.size(); ++i) {
        HCCL_DEBUG("[AllReduceAHC][RunInterAllReduce] run logicCard[%u] rank[%u] interRank=%u, interLinksSize=%u",
             logicCardList[i], rank, interRank, interLinksVector[i].size());
        CHK_RET(RunInstance(interRank, interLinksVector[i], interSlicesVector[i], tempAlg, HcclCMDType::HCCL_CMD_ALLREDUCE));
    }
    return HCCL_SUCCESS;
}

}